#!python3.9
# -*- coding: utf-8 -*-
# @Time    : 2025/2/17 10:54
# @Author  : fy
# @File    : pg_driver.py
# @Software: PyCharm
import os
import time
import logging
import psycopg2
import psycopg2.errors
import psycopg2.extensions
import psycopg2.extras
from threading import local


class PGDriver:
    """
    get conn url from env

    eg:
    export DATA_SOURCE_NAME="postgresql://user:password@ip:post/postgres"
    """
    __main_db_name__ = 'postgres'

    def __init__(self, url):
        self._url = url
        self.parsed_dsn = None
        self.initialized = False
        self._conn = local()

    def initialize(self, url):
        try:
            env_url = os.getenv('DATA_SOURCE_NAME')
            if env_url is not None:
                self._url = env_url
                url = env_url
            # Specify default schema is public.
            self.RealDictCursor = psycopg2.extras.RealDictCursor
            conn = psycopg2.connect(
                url, options="-c search_path=public",
                application_name='PG psycopg2'
            )
            # conn = psycopg2.connect(host='127.0.0.1', port=5237, user='psql', password='psql',
            #                                   database='postgres',
            #                                   connect_timeout=10, options="-c search_path=public",
            #     application_name='DBMind-openGauss-exporter'
            # )
            conn.cursor().execute('select 1;')
            conn.close()
            self._url = url
            self.parsed_dsn = psycopg2.extensions.parse_dsn(url)
            self.initialized = True
        except Exception as e:
            raise ConnectionError(e)

    @property
    def address(self):
        return '%s:%s' % (self.parsed_dsn['host'], self.parsed_dsn['port'])

    @property
    def dbname(self):
        return self.parsed_dsn['dbname']

    @property
    def username(self):
        return self.parsed_dsn['user']

    @property
    def pwd(self):
        return self.parsed_dsn['password']

    def query(self, stmt, timeout=0, return_tuples=False):
        if not self.initialized:
            raise AssertionError()
        cursor_dict = {}
        if not return_tuples:
            cursor_dict['cursor_factory'] = self.RealDictCursor
        try:
            conn = self.get_conn()
            with conn.cursor(
                    **cursor_dict
            ) as cursor:
                try:
                    start = time.monotonic()
                    if timeout > 0:
                        cursor.execute('SET statement_timeout = %d;' % (timeout * 1000))
                    cursor.execute(stmt)
                    result = cursor.fetchall()
                    logging.debug(len(result))
                except psycopg2.extensions.QueryCanceledError as e:
                    logging.error('%s: %s.' % (e.pgerror, stmt))
                    logging.info(
                        'Time elapsed during execution is %fs '
                        'but threshold is %fs.' % (time.monotonic() - start, timeout)
                    )
                    result = []
            conn.commit()
        except psycopg2.InternalError as e:
            logging.error("Cannot execute '%s' due to internal error: %s." % (stmt, e.pgerror))
            result = []
        except Exception as e:
            logging.exception(e)
            result = []
        return result

    def get_conn(self):
        """Cache the connection in the thread so that the thread can
        reuse this connection next time, thereby avoiding repeated creation.
        By this way, we can realize thread-safe database query,
        and at the same time, it can also have an ability similar to a connection pool. """
        # If query wants to connect to other database by force, we can generate and cache
        # the new connection as the following.

        db_name = self.dbname
        dsn = self._url

        if not hasattr(self._conn, db_name) or getattr(self._conn, db_name).closed:
            setattr(self._conn, db_name, psycopg2.connect(
                dsn, application_name='DBMind-openGauss-exporter'))
        # Check whether the connection is timeout or invalid.
        try:
            getattr(self._conn, db_name).cursor().execute('select 1;')
        except (
                psycopg2.InternalError,
                psycopg2.InterfaceError,
                psycopg2.errors.AdminShutdown,
                psycopg2.OperationalError
        ) as e:
            logging.warning(
                'Cached database connection to openGauss has been timeout due to %s, reconnecting.' % e
            )
            setattr(self._conn, db_name, psycopg2.connect(
                dsn, application_name='DBMind-openGauss-exporter'))
        except Exception as e:
            logging.error('Failed to connect to openGauss with cached connection (%s), try to reconnect.' % e)
            setattr(self._conn, db_name, psycopg2.connect(
                dsn, application_name='DBMind-openGauss-exporter'))
        return getattr(self._conn, db_name)
